from .base import BaseLoader
import os
from typing import Optional
from docx import Document
import PyPDF2

class DocumentLoader(BaseLoader):
    def __init__(self):
        """
        初始化文档加载器
        支持的格式：txt, pdf, docx
        """
        super().__init__()
        self.supported_formats = ['.txt', '.pdf', '.docx']
        self._pdf_reader = None
        self._docx_reader = None
    
    def load(self, file_path: str) -> str:
        """
        根据文件扩展名加载不同格式的文档
        参数:
            file_path: 文件路径
        返回:
            文件内容字符串
        异常:
            ValueError: 不支持的文件格式
            FileNotFoundError: 文件不存在
        """
        import os
        
        # 检查文件是否存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        # 获取文件扩展名
        file_ext = os.path.splitext(file_path)[1].lower()
        
        # 检查是否支持该文件格式
        if file_ext not in self.supported_formats:
            raise ValueError(f"不支持的文件格式: {file_ext}")
        
        # 根据文件格式调用相应的加载方法
        if file_ext == '.txt':
            return self._load_txt(file_path)
        elif file_ext == '.pdf':
            return self._load_pdf(file_path)
        elif file_ext == '.docx':
            return self._load_docx(file_path)
    
    def _load_txt(self, file_path: str) -> str:
        """加载TXT文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        except UnicodeDecodeError:
            with open(file_path, 'r', encoding='gbk') as f:
                return f.read()

    def _load_pdf(self, file_path: str) -> str:
        """加载PDF文件"""
        try:
            if self._pdf_reader is None:
                from PyPDF2 import PdfReader
                self._pdf_reader = PdfReader
            
            reader = self._pdf_reader(file_path)
            text = []
            for page in reader.pages:
                text.append(page.extract_text())
            return '\n'.join(text)
        except ImportError:
            raise ImportError("请安装 PyPDF2: pip install PyPDF2")

    def _load_docx(self, file_path: str) -> str:
        """加载DOCX文件"""
        try:
            if self._docx_reader is None:
                from docx import Document
                self._docx_reader = Document
            
            doc = self._docx_reader(file_path)
            text = []
            for para in doc.paragraphs:
                if para.text.strip():
                    text.append(para.text)
            return '\n'.join(text)
        except ImportError:
            raise ImportError("请安装 python-docx: pip install python-docx")

    def load_document(self, file_path: str) -> Optional[str]:
        """
        加载文档内容
        
        参数:
            file_path: 文件路径
            
        返回:
            文档内容字符串，如果加载失败则返回None
        """
        try:
            file_extension = os.path.splitext(file_path)[1].lower()
            
            if file_extension == '.txt':
                return self._load_txt(file_path)
            elif file_extension == '.pdf':
                return self._load_pdf(file_path)
            elif file_extension == '.docx':
                return self._load_docx(file_path)
            else:
                print(f"不支持的文件格式: {file_extension}")
                return None
                
        except Exception as e:
            print(f"加载文件出错 {file_path}: {str(e)}")
            return None

